import time,os.path
from threading import Thread
from keras.callbacks import Callback
from keras.layers import Input, Dense
from keras.models import Model
from keras.optimizers import Adam
from PyQt5.QtGui import QIntValidator,QDoubleValidator,QTextCursor
from PyQt5.QtWidgets import QWidget,QFileDialog
from PyQt5.QtCore import pyqtSignal
from keras.models import load_model
from gui.windows.autoencoder import Ui_autoencoder
from config.network import custom_comp

class AutoencoderWindow(QWidget, Ui_autoencoder):
    signal_update_text = pyqtSignal(str,bool)
    signal_finish_train = pyqtSignal()

    def __init__(self,startWindow):
        super(AutoencoderWindow, self).__init__()
        self.setupUi(self)
        self.setFixedSize(646, 774)
        self.startWindow=startWindow
        self.param_dict=startWindow.param_dict
        self.dataWindow=startWindow.dataWindow
        self.initial_property()
        self.initial_widget()
        self.connect_slot()

    def initial_property(self):
        self.text_string = ''
        self.traces = None
        self.autoencoder = None
        self.stop_train=False
        self.network_type={
            'MLP':self.Autoencoder_MLP,
            'CNN':self.Autoencoder_CNN,
            'LSTM':self.Autoencoder_LSTM,
        }

    def initial_widget(self):
        self.network.addItems(self.network_type)
        self.network.setCurrentText(list(self.network_type.keys())[0])
        self.epochs.setText(str(300))
        self.batch_size.setText(str(128))
        self.learning_rate.setText(str(0.0001))
        self.target_dim.setText(str(512))
        self.set_func_choose_status()
        self.set_validator()
        with open("resource/style.qss", encoding='UTF-8') as f:
            self.setStyleSheet(f.read())

    def connect_slot(self):
        self.gen_model.clicked.connect(self.set_func_choose_status)
        self.redim_traces.clicked.connect(self.set_func_choose_status)
        self.autoencoder_path_browser.clicked.connect(self.browse_autoencoder_path)
        self.start_run.clicked.connect(self.start_process)
        self.signal_update_text.connect(self.update_text)
        self.signal_finish_train.connect(lambda :self.save_autoencoder() or self.start_run.setText("开始执行"))
        self.related_info.textChanged.connect(lambda: self.related_info.moveCursor(QTextCursor.End))

    def start_process(self):
        try:
            if self.start_run.text()=='结束训练':
                self.stop_train=True
                return
            if self.gen_model.isChecked():
                self.train_autoencoder()
            else:
                if self.autoencoder is None:
                    raise ValueError("请先加载或训练自编码器模型！")
                encoder = Model(inputs=self.autoencoder.inputs, outputs=self.autoencoder.get_layer('encoder_output').output)
                input_dim, output_dim = encoder.input_shape[1],encoder.output_shape[1]
                if input_dim != self.traces.shape[1]:
                    if output_dim==self.traces.shape[1]:
                        raise ValueError('能量曲线的点数与模型输入维度不一致（可能已对曲线降维）')
                    raise ValueError('能量曲线的点数与模型输入维度不一致（可能选择了错误的模型）')
                self.traces = encoder.predict(self.traces,verbose=False)
                self.dataWindow.traces=self.traces
                self.dataWindow.update_traces_show(True)
                self.update_text('降维完毕.')
        except Exception as e:
            self.update_text(repr(e))
            self.update_text('处理失败！')

    def train_autoencoder(self):
        epochs = int(self.epochs.text())
        batch_size = int(self.batch_size.text())
        input_dim = self.traces.shape[1]
        learning_rate = float(self.learning_rate.text())
        target_dim = int(self.target_dim.text())
        autoencoder_func = self.network_type[self.network.currentText()]
        self.autoencoder = autoencoder_func(input_dim,target_dim, learning_rate)
        th_train = Thread(target=self.train_th, args=(epochs, batch_size))
        self.start_run.setText('结束训练')
        th_train.start()

    def train_th(self,epochs,batch_size):
        class StopTrainingCallback(Callback):
            def __init__(self,parent):
                super(StopTrainingCallback,self).__init__()
                self.parent=parent
                self.model=parent.autoencoder
            def on_epoch_end(self, epoch, logs=None):
                logs = logs or {}
                message = f"Epoch {epoch + 1}\t"
                for key, value in logs.items():
                    message += f"{key}: {value:.4f}"
                self.parent.signal_update_text.emit(message,False)
                if self.parent.stop_train:
                    self.model.stop_training = True
        self.signal_update_text.emit("开始训练...", False)
        self.autoencoder.fit(self.traces, self.traces, epochs=epochs, batch_size=batch_size, shuffle=True,callbacks=[StopTrainingCallback(self)],verbose=False)
        self.signal_update_text.emit("训练完成.", False)
        self.signal_finish_train.emit()

    def save_autoencoder(self):
        try:
            ctime = time.localtime()
            dir = self.param_dict['model_file_dir'] + f"{ctime.tm_year}-{ctime.tm_mon}-{ctime.tm_mday}/"
            if os.path.exists(dir) is False:
                os.makedirs(dir)
            hms = f'{ctime.tm_hour:02d}{ctime.tm_min:02d}{ctime.tm_sec:02d}'
            file_name = f"Autoencoder_{self.network.currentText()}_{hms}.h5"
            fname, _ = QFileDialog.getSaveFileName(self, "Save file", dir + file_name, "*.h5")
            if fname:
                self.autoencoder.save(fname)
                self.signal_update_text.emit(f"自编码器模型已保存到{fname}.", False)
                self.autoencoder_path.setText(fname)
            else:
                self.signal_update_text.emit("模型未保存！", False)
        except Exception as e:
            self.signal_update_text.emit(repr(e), False)
            self.signal_update_text.emit("模型保存失败！", False)

    def browse_autoencoder_path(self):
        try:
            fname, _ = QFileDialog.getOpenFileName(self, "Open file", self.param_dict['model_file_dir'], "*.h5")
            if fname:
                self.autoencoder = load_model(fname, custom_objects=custom_comp)
                input_dim = self.autoencoder.get_layer(index=0).input_shape[0]
                if input_dim[1]!=self.traces.shape[1]:
                    self.autoencoder=None
                    raise RuntimeError("自编码器模型与能量曲线不匹配！")
                self.autoencoder_path.setText(fname)
                self.update_text("加载自编码器模型成功.")
        except Exception as e:
            self.update_text(repr(e))
            self.update_text("加载自编码模型失败(可能选择了错误的文件)！")

    def set_validator(self):
        int_validator1 = QIntValidator(10, 100000, self)
        int_validator2 = QIntValidator(10, 10000, self)
        double_validator=QDoubleValidator(0,1,6,self)
        double_validator.setNotation(QDoubleValidator.StandardNotation)
        self.epochs.setValidator(int_validator1)
        self.batch_size.setValidator(int_validator2)
        self.target_dim.setValidator(int_validator2)
        self.learning_rate.setValidator(double_validator)

    def set_func_choose_status(self):
        generate_model = self.gen_model.isChecked()
        redim_traces = self.redim_traces.isChecked()
        self.autoencoder_path_L.setHidden(generate_model)
        self.autoencoder_path.setHidden(generate_model)
        self.autoencoder_path_browser.setHidden(generate_model)
        self.network_L.setHidden(redim_traces)
        self.network.setHidden(redim_traces)
        self.epochs_L.setHidden(redim_traces)
        self.epochs.setHidden(redim_traces)
        self.batch_size_L.setHidden(redim_traces)
        self.batch_size.setHidden(redim_traces)
        self.learning_rate_L.setHidden(redim_traces)
        self.learning_rate.setHidden(redim_traces)
        self.target_dim.setHidden(redim_traces)
        self.target_dim_L.setHidden(redim_traces)

    def Autoencoder_MLP(self,input_dim,encoding_dim=512,learning_rate=0.0001):
        input = Input(shape=(input_dim,))
        encode = Dense(2048, activation='relu')(input)
        encode = Dense(1024, activation='relu')(encode)
        encode = Dense(encoding_dim,name='encoder_output')(encode)
        decode = Dense(1024, activation='relu')(encode)
        decode = Dense(2048, activation='relu')(decode)
        decoder_output = Dense(input_dim, activation='tanh')(decode)
        Autoencoder = Model(inputs=input, outputs=decoder_output)
        Autoencoder.compile(optimizer=Adam(learning_rate),loss='mse',metrics=['accuracy'])
        return Autoencoder

    def Autoencoder_CNN(self,input_dim,encoding_dim=512,learning_rate=0.0001):
        input = Input(shape=(input_dim,))
        encode = Dense(2048, activation='relu')(input)
        encode = Dense(1024, activation='relu')(encode)
        encode = Dense(encoding_dim,name='encoder_output')(encode)
        decode = Dense(1024, activation='relu')(encode)
        decode = Dense(2048, activation='relu')(decode)
        decoder_output = Dense(input_dim, activation='tanh')(decode)
        Autoencoder = Model(inputs=input, outputs=decoder_output)
        Autoencoder.compile(optimizer=Adam(learning_rate),loss='mse',metrics=['accuracy'])
        return Autoencoder

    def Autoencoder_LSTM(self,input_dim,encoding_dim=512,learning_rate=0.0001):
        input = Input(shape=(input_dim,))
        encode = Dense(2048, activation='relu')(input)
        encode = Dense(1024, activation='relu')(encode)
        encode = Dense(encoding_dim,name='encoder_output')(encode)
        decode = Dense(1024, activation='relu')(encode)
        decode = Dense(2048, activation='relu')(decode)
        decoder_output = Dense(input_dim, activation='tanh')(decode)
        Autoencoder = Model(inputs=input, outputs=decoder_output)
        Autoencoder.compile(optimizer=Adam(learning_rate),loss='mse',metrics=['accuracy'])
        return Autoencoder

    def update_text(self, data, override=False):
        if override:
            self.text_string = ''
        if type(data) is dict:
            self.text_string += 'trs file head:\n'
            for k in data:
                self.text_string += f"  {k}: {data[k]}\n"
            self.text_string += '\n'
        else:
            if data[-1] == '#':
                self.text_string += str(data)[:-1]
            elif len(data)>3 and data[-3:]=='...':
                self.text_string += "(Please Wait)" + str(data) + '\n'
            elif data[-1]=='.':
                self.text_string += "(Execution Success)"+str(data) + '\n'
            elif data[-1]=='!' or data[-1]=='！':
                self.text_string += "(Execution Fail)" + str(data) + '\n'
            else:
                self.text_string += str(data) + '\n'
        self.related_info.setText(self.text_string)
        if len(self.text_string) > 10240:
            self.text_string = self.text_string[7680:]

    def showEvent(self, a0):
        self.traces=self.dataWindow.traces
        self.text_string=''
        self.related_info.setText('')

    def closeEvent(self, a0):
        self.text_string = ''
        self.update_text("#", True)
        self.startWindow.dataWindow.update_process_tips('自编码器窗口已关闭.')
        self.startWindow.dataWindow.show()